#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test results')
parser.add_argument('-t', '--time-step', type=int, help='Time step to plot', default=-1)
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-v', '--var', help='Variable to plot', default='v')
args = parser.parse_args()

f = xr.open_dataset('swm_cp.360x180.h0.nc')

fig = plt.figure(figsize=(8, 8))
# proj = ccrs.NorthPolarStereo(central_longitude=0)
proj = ccrs.PlateCarree()
ax = fig.add_subplot(1, 1, 1, projection=proj)
# ax.set_extent([-180, 180, 30, 90], ccrs.PlateCarree())

if args.var == 'u':
  levels = np.linspace(-24, 24, 13)
  colors = ListedColormap([
    [0.01, 0.12, 0.80],
    [0.11, 0.31, 0.98],
    [0.38, 0.74, 0.98],
    [0.47, 0.84, 0.99],
    [0.56, 0.93, 0.99],
    [0.78, 1.00, 1.00],
    [1.00, 0.98, 0.33],
    [1.00, 0.73, 0.18],
    [1.00, 0.51, 0.15],
    [1.00, 0.21, 0.11],
    [0.82, 0.00, 0.08],
    [0.59, 0.00, 0.06]
  ])
  plot_contour_map(ax, f.u[args.time_step,0,:,:], cmap=colors, with_grid=True, left_string='U (m/s)', levels=levels, show_lat_labels=False)
if args.var == 'v':
  levels = np.linspace(-15, 15, 13)
  colors = ListedColormap([
    [0.01, 0.12, 0.80],
    [0.11, 0.31, 0.98],
    [0.38, 0.74, 0.98],
    [0.47, 0.84, 0.99],
    [0.56, 0.93, 0.99],
    [0.78, 1.00, 1.00],
    [1.00, 0.98, 0.33],
    [1.00, 0.73, 0.18],
    [1.00, 0.51, 0.15],
    [1.00, 0.21, 0.11],
    [0.82, 0.00, 0.08],
    [0.59, 0.00, 0.06]
  ])
  plot_contour_map(ax, f.v[args.time_step,0,:,:], cmap=colors, with_grid=True, left_string='V (m/s)', levels=levels, show_lat_labels=False)

plt.tight_layout()
if args.output:
  plt.savefig(args.output, format='pdf')
else:
  plt.show()
